from migration.base_migrator import BaseMigrator
from utils.common import CommonUtils
import os

class TableMigrator(BaseMigrator):
    def __init__(self, mysql_db, kingbase_db, migration_config, logger, task_name):
        super().__init__(mysql_db, kingbase_db, migration_config, logger, task_name)
        
    def get_tables_to_migrate(self):
        """
        获取需要迁移的表
        """
        all_tables = self.mysql_db.get_all_tables()
        include_tables = self.migration_config.get('include_tables', [])
        exclude_tables = self.migration_config.get('exclude_tables', [])
        
        # 如果指定了包含表，则只迁移包含的表
        if include_tables:
            tables_to_migrate = [table for table in all_tables if table in include_tables]
        else:
            tables_to_migrate = all_tables
        
        # 排除指定的表
        if exclude_tables:
            tables_to_migrate = [table for table in tables_to_migrate if table not in exclude_tables]
        
        self.logger.info(f"共需要迁移 {len(tables_to_migrate)} 个表")
        return tables_to_migrate
    
    def migrate_table_structure(self, table_name):
        """
        迁移单个表的结构
        """
        self.logger.info(f"开始迁移表结构: {table_name}")
        try:
            # 检查目标数据库是否已存在该表
            # 使用小写的模式名，与kingbase_db.py保持一致
            target_schema_lower = self.target_schema.lower()
            if self.kingbase_db.table_exists(target_schema_lower, table_name):
                if self.migration_config.get('skip_existing_tables', False):
                    self.logger.warning(f"表 {target_schema_lower}.{table_name} 已存在，跳过")
                    return None, None
                else:
                    self.logger.info(f"表 {target_schema_lower}.{table_name} 已存在，删除后重新创建")
                    self.kingbase_db.drop_table(target_schema_lower, table_name)
            
            # 获取表结构
            table_structure = self.mysql_db.get_table_structure(table_name)
            
            # 生成创建表的SQL和列注释
            create_table_sql, column_comments = self._generate_create_table_sql(table_name, table_structure)
            
            # 记录生成的SQL语句，用于调试
            self.logger.debug(f"生成的CREATE TABLE语句:\n{create_table_sql}")
            self.logger.debug(f"生成的列注释数量: {len(column_comments)}")
            if column_comments:
                self.logger.debug(f"第一个列注释语句:\n{column_comments[0]}")
            
            # 对特定表添加更详细的调试信息
            if table_name == 'AgencyProcessingLog':
                self.logger.info(f"正在处理表: {table_name}")
                self.logger.info(f"表结构包含的列数: {len(table_structure)}")
                # 显示所有列名和KEY
                for column in table_structure:
                    self.logger.info(f"MySQL列名: {column['COLUMN_NAME']}, KEY: {column['COLUMN_KEY']}")
                # 显示完整的CREATE TABLE语句
                self.logger.info(f"完整的CREATE TABLE语句:\n{create_table_sql}")
                # 显示完整的列注释语句
                if column_comments:
                    self.logger.info(f"完整的列注释语句:\n{column_comments}")
            
            # 执行创建表SQL
            # 使用小写的模式名，与kingbase_db.py保持一致
            target_schema_lower = self.target_schema.lower()
            self.kingbase_db.create_table(target_schema_lower, create_table_sql)
            
            # 执行列注释SQL
            for comment_sql in column_comments:
                self.kingbase_db.execute_update(comment_sql)

            # 迁移表的注释
            table_comment = self.mysql_db.get_table_comment(table_name)
            if table_comment:
                self._add_table_comment(table_name, table_comment)
            
            # 检查是否有LastEditTime字段，如果有则创建更新时间触发器
            has_last_edit_time = any(column['COLUMN_NAME'].lower() == 'lastedittime' for column in table_structure)
            if has_last_edit_time:
                self._create_last_edit_time_trigger(table_name)
            
            self.logger.info(f"表结构迁移成功: {target_schema_lower}.{table_name}")
            return create_table_sql, column_comments
        except Exception as e:
            self.logger.error(f"表结构迁移失败: {table_name}\n{e}")
            raise
    
    def _generate_create_table_sql(self, table_name, table_structure):
        """
        生成创建表的SQL语句，确保兼容PostgreSQL/KingbaseES语法
        """
        columns_sql = []
        primary_keys = []
        indexes = []
        foreign_keys = []
        comments = []
        column_comments = []
        
        # 表级别设置
        table_options = []
        
        for column in table_structure:
            column_name = column['COLUMN_NAME']
            column_type = column['COLUMN_TYPE'].lower()
            is_nullable = column['IS_NULLABLE']
            column_key = column['COLUMN_KEY']
            column_default = column['COLUMN_DEFAULT']
            extra = column['EXTRA']
            column_comment = column.get('COLUMN_COMMENT', '')
            character_set = column.get('CHARACTER_SET_NAME')
            collation = column.get('COLLATION_NAME')
            
            # 转换数据类型
            converted_type = self._convert_data_type(column_type)
            
            # 额外检查：确保整数类型没有长度参数
            # 这是一个额外的安全保障，确保integer(4)这样的语法不会出现在最终的SQL中
            if any(keyword in converted_type.lower() for keyword in ['int', 'integer', 'smallint', 'mediumint', 'bigint', 'tinyint']):
                # 移除可能存在的任何参数部分
                base_type = converted_type.split('(')[0].strip()
                converted_type = base_type
            
            # 构建列定义，所有列名都用双引号括起来以保持大小写
            column_def = f'"{column_name}" {converted_type}'
            
            # 处理NOT NULL
            if is_nullable == 'NO':
                # 根据转换说明，VARCHAR类型不需要NOT NULL
                if not any(varchar_type in converted_type.lower() for varchar_type in ['varchar', 'char', 'text']):
                    column_def += " NOT NULL"
            
            # 处理默认值
            if column_default is not None:
                # 特殊处理默认值为0000-00-00 00:00:00的情况
                if isinstance(column_default, str) and column_default.lower() in ['0000-00-00 00:00:00', '0000-00-00']:
                    if 'timestamp' in converted_type.lower():
                        column_default = '1000-01-01 00:00:00'
                    elif 'date' in converted_type.lower():
                        column_default = '1000-01-01'
                column_def = self._handle_default_value(column_def, column_default, converted_type)
            
            # 处理自增
            if 'auto_increment' in extra.lower():
                # 根据转换说明，AUTO_INCREMENT字段应该转换为SERIAL PRIMARY KEY
                if column_key == 'PRI':
                    # 对于主键自增字段，使用SERIAL PRIMARY KEY，并为列名添加双引号以保持大小写
                    column_def = f'"{column_name}" SERIAL PRIMARY KEY'
                else:
                    # 对于非主键自增字段，使用GENERATED BY DEFAULT AS IDENTITY
                    column_def += " GENERATED BY DEFAULT AS IDENTITY"
            
            # 处理字符集和排序规则
            if character_set and collation:
                # PostgreSQL/KingbaseES在列级别设置字符集的方式不同
                # 我们可以添加注释来说明原始设置
                comments.append(f"-- 列 {column_name} 原始字符集: {character_set}, 排序规则: {collation}")
            
            # 记录列注释，将在表创建后添加
            if column_comment:
                # 所有列名都用双引号括起来以保持大小写
                formatted_column_name = f'"{column_name}"'
                # 转义单引号
                escaped_comment = column_comment.replace("'", "''")
                # 使用小写的模式名，与kingbase_db.py保持一致
                target_schema_lower = self.target_schema.lower()
                column_comments.append(f"COMMENT ON COLUMN {target_schema_lower}.\"{table_name}\".{formatted_column_name} IS '{escaped_comment}';")
            
            columns_sql.append(column_def)
            
            # 记录主键
            if column_key == 'PRI':
                # 检查列定义是否已包含SERIAL PRIMARY KEY，避免重复添加
                if not ('auto_increment' in extra.lower() and 'SERIAL PRIMARY KEY' in column_def):
                    # 检查是否已经有自增主键
                    has_serial_primary_key = any('SERIAL PRIMARY KEY' in col for col in columns_sql)
                    if not has_serial_primary_key:
                        primary_keys.append(column_name)
            # 记录唯一索引（在PostgreSQL中作为表约束处理）
            elif column_key == 'UNI':
                indexes.append(f'CONSTRAINT idx_{table_name}_{column_name}_unique UNIQUE ("{column_name}")')
            # 记录普通索引（这些将在表创建后单独创建）
            elif column_key == 'MUL':
                # 我们不会在CREATE TABLE语句中包含普通索引，因为PostgreSQL语法不同
                # 这些索引将在单独的SQL文件中生成或在表创建后单独执行
                pass
        
        # 构建主键定义（排除已经被设置为SERIAL PRIMARY KEY的字段）
        primary_key_sql = ""
        serial_primary_keys = []
        # 找出已经被设置为SERIAL PRIMARY KEY的字段
        for column in table_structure:
            if 'auto_increment' in column['EXTRA'].lower() and column['COLUMN_KEY'] == 'PRI':
                serial_primary_keys.append(column['COLUMN_NAME'])
        # 只为主键列表中不在serial_primary_keys中的字段创建主键约束
        non_serial_primary_keys = [pk for pk in primary_keys if pk not in serial_primary_keys]
        
        if non_serial_primary_keys:
            # 主键字段用双引号括起来以保持大小写
            quoted_pk_columns = [f'"{pk}"' for pk in non_serial_primary_keys]
            primary_key_sql = f",\n    PRIMARY KEY ({', '.join(quoted_pk_columns)})"
        
        # 构建索引SQL
        indexes_sql = ""
        if indexes:
            indexes_sql = f",\n    " + f",\n    ".join(indexes)
        
        # 构建表选项
        table_options_sql = """
WITH (
    OIDS = FALSE
)"""
        
        # 构建完整的CREATE TABLE语句
        columns_str = ',\n    '.join(columns_sql)
        
        # 确保不会有多余的逗号
        if primary_key_sql:
            columns_str += primary_key_sql
        if indexes_sql:
            columns_str += indexes_sql
        
        # 使用小写的模式名，与kingbase_db.py保持一致
        target_schema_lower = self.target_schema.lower()
        create_table_sql = "\n".join(comments) + f'''

CREATE TABLE IF NOT EXISTS {target_schema_lower}."{table_name}" (
    {columns_str}
){table_options_sql};
'''
        
        return create_table_sql, column_comments
    
    def _add_table_comment(self, table_name, comment):
        """
        为表添加注释
        """
        try:
            # 转义单引号
            escaped_comment = comment.replace("'", "''")
            # 使用小写的模式名，与kingbase_db.py保持一致
            target_schema_lower = self.target_schema.lower()
            sql = f"COMMENT ON TABLE {target_schema_lower}.\"{table_name}\" IS '{escaped_comment}';"
            self.kingbase_db.execute_update(sql)
            self.logger.info(f"表注释添加成功: {target_schema_lower}.{table_name}")
        except Exception as e:
            self.logger.warning(f"表注释添加失败: {table_name}\n{e}")
    
    def _create_last_edit_time_trigger(self, table_name):
        """
        根据转换说明，为表创建LastEditTime字段的更新时间触发器
        """
        try:
            # 使用小写的模式名，与kingbase_db.py保持一致
            target_schema_lower = self.target_schema.lower()
            
            # 创建触发器函数
            trigger_func_name = f"update_timestamp"
            trigger_func_sql = f"""
CREATE OR REPLACE FUNCTION {target_schema_lower}.{trigger_func_name}()
RETURNS TRIGGER AS $$
BEGIN
    NEW."LastEditTime" = CURRENT_TIMESTAMP;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;
"""
            self.kingbase_db.execute_update(trigger_func_sql)
            
            # 创建触发器
            trigger_name = f"set_LastEditTime"
            trigger_sql = f'''
CREATE TRIGGER {trigger_name}
BEFORE UPDATE ON {target_schema_lower}."{table_name}"
FOR EACH ROW
EXECUTE PROCEDURE {target_schema_lower}.{trigger_func_name}();
'''
            self.kingbase_db.execute_update(trigger_sql)
            self.logger.info(f"更新时间触发器创建成功: {target_schema_lower}.{table_name}")
        except Exception as e:
            self.logger.warning(f"更新时间触发器创建失败: {table_name}\n{e}")
    
    def create_indexes(self, table_name, indexes):
        """
        为表创建索引
        """
        for index in indexes:
            try:
                # 使用小写的模式名
                target_schema_lower = self.target_schema.lower()
                
                # 构建索引列定义（用双引号括起来以保持大小写）
                quoted_columns = [f'"{col}"' for col in index['columns']]
                columns_str = ', '.join(quoted_columns)
                
                # 构建索引类型关键字
                index_type = 'UNIQUE' if index['is_unique'] else ''
                
                # 构建索引名（确保唯一性）
                # 替换表名中的连字符为下划线，因为PostgreSQL/KingbaseES不允许索引名中包含连字符
                safe_table_name = table_name.replace('-', '_')
                index_name = f"idx_{safe_table_name}_{'_'.join(index['columns']).lower()}"
                
                # 构建创建索引的SQL
                create_index_sql = f'CREATE {index_type} INDEX IF NOT EXISTS {index_name} ON {target_schema_lower}."{table_name}" ({columns_str});'
                
                # 执行SQL
                self.kingbase_db.execute_update(create_index_sql)
                self.logger.info(f"索引创建成功: {index_name}")
            except Exception as e:
                self.logger.warning(f"索引创建失败: {index['name']}\n{e}")
    
    def save_create_table_sqls(self, create_table_sqls):
        """
        保存创建表的SQL到文件
        """
        # 保存创建表的SQL
        tables_sql_file = os.path.join(self.reports_dir, 'tables_create.sql')
        CommonUtils.write_text_file(tables_sql_file, '\n\n'.join(create_table_sqls))
        self.logger.info(f"表创建SQL保存成功: {tables_sql_file}")
        
        # 保存创建模式的SQL
        schema_sql_file = os.path.join(self.reports_dir, 'schema_create.sql')
        schema_sql = f"CREATE SCHEMA IF NOT EXISTS {self.target_schema.lower()};"
        CommonUtils.write_text_file(schema_sql_file, schema_sql)
        self.logger.info(f"模式创建SQL保存成功: {schema_sql_file}")